import time
import signal
import sys
import socket
import threading
import psutil
import multiprocessing
from server import Server
from utils import *
from common import *

# ==================== 从服务端 ====================
class SlaveServer:
    def __init__(self, master_host, master_port, client_id):
        self.master_host = master_host
        self.master_port = master_port
        self.client_id = client_id
        self.conn = None
        self.running = threading.Event()
        # 存储主机信息
        self.host_info = HostInfo(client_id)
        self.host_info_lock = threading.RLock()
        
    def start(self):
        """启动从服务端"""
        self.running.set()
        self._handle()
        self.stop()
        
    def _is_able_receive_task(self):
        """根据主机当前状态判断是否可以接收任务"""
        self._update_self_host_info()
        with self.host_info_lock:
            return False if self.host_info.on_task else True  # 目前策略
    
    def _update_self_host_info(self):
        memory_info = psutil.virtual_memory()
        def query_task_status(task_id, pid):
            return {'path':{'aaaaaa':"aaa", "bbbb": "bbb"}} 
        with self.host_info_lock:
            self.host_info.total_memory = '{:.2f}'.format(memory_info.total/(1024**3))
            self.host_info.available_memory = '{:.2f}'.format(memory_info.available / (1024**3))
            self.host_info.used_memory = '{:.2f}'.format(memory_info.used / (1024**3))
            self.host_info.mem_percent = f'{memory_info.percent}%'
            # TODO 检查当前是否还有任务执行，更新当前主机的任务信息
            self.host_info.on_task = False
            for task_info in self.host_info.tasks_info:
                if psutil.pid_exists(task_info.PID):
                    task_info.status = TaskStatus.Running
                    self.host_info.on_task = True
                else:
                    task_info.status = TaskStatus.Finished
                    self.host_info.on_task_num -= 1
    
    def _handle_new_tasks(self, task: Task):
        """包含任务解析,  执行任务返回结果"""
        eval_list = task.convert_eval_task()
        logger.debug(f"[{self.client_id}] eval_list: {eval_list}")
        pid = self._execute_task(eval_list[0])
        if pid:
            return True, pid
        else:
            return False, pid
    
    def _core(self, serve: Server):        
        # 任务处理循环
        while serve.running.is_set() and serve.registered.is_set():
            response = serve.get_message()
            # 路由
            if response['type'] == MessageType.NewTasks:
                if self._is_able_receive_task():
                    serve.send_message(build_get_task())
                    logger.info(f"[{self.client_id} Serve] send get task message")
                else:
                    logger.debug(f"Received a new task notification, but the current status is unacceptable for the task.")

            elif response['type'] == MessageType.Task:
                # 没有任务休眠
                if response['task']:
                    task = Task.model_validate(response['task'])
                    # 有任务执行（处理主要的任务逻辑函数）
                    try:
                        result, pid = self._handle_new_tasks(task)
                        # TODO 如果失败需要删除多余文件夹
                        if result:
                            serve.send_message(build_allow())
                            # 根据任务构建一个任务信息
                            task_info_list = self._get_task_info_by_(task, pid)
                            with self.host_info_lock:
                                # 将任务信息加入到表中
                                self.host_info.set_task_info(task_info_list)
                        else:
                            serve.send_message(build_refuse(f"执行评测任务失败"))
                    except Exception as e:
                        logger.error(f"{e}")
                        serve.send_message(build_refuse(f"{e}"))
                    # 完成后再发送一次请求
                    if self._is_able_receive_task(): 
                        serve.send_message(build_get_task())
                    else:
                        logger.debug(f"There is an existing task. No need to send the message for obtaining the task.")

            elif response['type'] == MessageType.Exit:
                serve.stop()
    def _handle_host_info(self, serve: Server):
        while self.running.is_set() and serve.registered.is_set():
            try:
                # 收到hostinfo
                serve.get_message()
                logger.debug(f"[{serve.name}] recv slave_info msg")
                # 先更新自身信息
                self._update_self_host_info()
                # 发送序列化后的信息
                with self.host_info_lock:
                    host_info = self.host_info.serialize()
                logger.debug(f"[{host_info}] send host_info msg")
                serve.send_message(build_slave_info(host_info))
            except:
                logger.error(f"Heartbeats occur exception")
                break

    def _handle(self):
        """主循环"""
        with socket.socket() as s:
            linger = struct.pack('ii', 1, 5)  # 启用 SO_LINGER，延迟时间为 10 秒
            s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, linger)
            s.connect((self.master_host, self.master_port))
            serve = Server(s, self.master_host, self.master_port, remote_name=self.client_id+'-master', self_name=self.client_id)
            serve.set_processor({
                self._core: [MessageType.NewTasks, MessageType.Task, MessageType.Refuse, MessageType.Allow],
                self._send_heartbeats: [],
                self._handle_host_info: [MessageType.SlaveInfo],
                self._handle_putong_putong: [MessageType.PutongPutong],
            })
            
            if self._register(s): serve.registered.set()

            if self.running.is_set(): serve.start()

    def _register(self, conn: socket.socket) -> bool:
        MessageProtocol.send(conn, build_register(self.client_id))
        msg = MessageProtocol.recv(conn)
        if msg.get('type') == MessageType.Refuse:
            logger.info(f"{msg['message']}")
            self.stop()
            return False
        # 更新自身主机信息
        self._update_self_host_info()
        return True

    # 心跳处理
    def _send_heartbeats(self, serve: Server):
        """发送心跳"""
        while self.running.is_set() and serve.registered.is_set():
            try:
                self._update_self_host_info()
                with self.host_info_lock:
                    serve.send_message(build_heartbeat(self.client_id, self.host_info.serialize()))
                time.sleep(HEARTBEAT_PER)
            except:
                logger.error(f"Heartbeats occur exception")
                break

    def _execute_task_test(self, task: Task) -> bool:
        """执行任务（示例）"""
        print(f"Executing task {task.id}")
        return True  # 模拟成功
    
    def _execute_task(self, task: EvalTask):
        """执行成功返回PID否则None"""
        try:
            # 生成配置文件
            # config_path = generate_config_file(task)
            # 根据配置文件执行任务
            p = multiprocessing.Process(target=test_task, args=(task,))
            p.daemon = True
            p.start()
            # 报错任务信息
            # TODO 应该加锁
            pid = p.pid
            # task_list.update(task_info)
            # 返回状态
            return pid
        except Exception as e:
            # if 'p' in locals() and p.is_alive():
            #     p.terminate()
            #     # TODO 清理子进程产生的垃圾
            #     p.join()  # 确保子进程已经终止
            return None
    
    def _get_task_info_by_(self, task: Task, pid) -> list[TaskInfo]:
        task_info = []
        for dataset_name, subset_name in task.dataset.items():
            task_info.append(TaskInfo(
                id=task.id,
                dataset_name=dataset_name,
                model=task.model,
                PID=pid,
                subset_list=subset_name,
                status=TaskStatus.Running
            ))
        return task_info
    
    def _handle_putong_putong(self, serve: Server):
        while serve.running.is_set() and serve.registered.is_set():
            logger.debug(f"putong putong")
            msg = serve.get_message()
            logger.debug(f"putong putong !!!")
            data = [TaskInfo.deserialize(t) for t in msg['data']]
            if len(data) != 0:
                with self.host_info_lock:
                    removed_task_set = {(task.id, task.dataset_name) for task in data}

                    self.host_info.tasks_info = [
                        task_info for task_info in self.host_info.tasks_info
                        if (task_info.id, task_info.dataset_name) not in removed_task_set
                    ]
            if self._is_able_receive_task():
                serve.send_message(build_get_task())
            
    
    def stop(self):
        self.running.clear()
        # MessageProtocol.send(self.conn, {'type': MessageType.Exit})
        print("host exit!")

def signal_handler(sig, frame):
    print("Interrupt signal received. Stopping servers...")
    if slave: slave.stop()
    sys.exit(0)
    
def slave_start():
    # 启动从服务端
    global slave
    slave = SlaveServer('localhost', 9000, 'slave_01')
    slave.start()
    
def test_task(task: EvalTask):
    logger.info(f'{task} 任务正在执行')
    time.sleep(6)
    logger.info(f"执行完成")
    
# ==================== 使用示例 ====================
if __name__ == '__main__':
    # 注册信号处理函数
    signal.signal(signal.SIGINT, signal_handler)
    slave_start()